import xml.sax.handler
import sqlite3
import sys
import logging
'''
Created on 15 Mar 2010

@author: A Hopkins
'''

class nmapXMLHandler(xml.sax.handler.ContentHandler):
    '''
    classdocs
    '''
    def __init__(self, dbOutputFile):
        '''
        Constructor
        '''
        self.inAddr = None
        self.inPortId = None
        self.inProtocol = None

        try: 
            self.conn = sqlite3.connect(dbOutputFile)  
            self.cursor = self.conn.cursor();
        except sqlite3.OperationalError as strerror:
            logging.error ("SQLITE error connecting to database: " + str(strerror))  
            sys.exit(2)  
        
    
    # used for parsing the XML file, calls the relevant methods to insert items
    def startElement(self, name, attributes):   
        if name == "address":
            if (attributes.get("addrtype","") == "ipv4"):
                self.inIPaddr = attributes.get("addr","")
                self.inAddrType = attributes.get("addrtype","")
                self.insertNewHost(self.inIPaddr)
            elif (attributes.get("addrtype","") == "mac"):
                inMacAddr = attributes.get("addr","")
                self.updateHostEntry("macaddress", inMacAddr)
        elif name == "port":  
            self.inPortId = attributes.get("portid","")
            self.inProtocol = attributes.get("protocol","")
            inProduct = attributes.get("product","")
            self.insertNewService(self.inIPaddr, self.inProtocol, self.inPortId)       
        elif name == "service":
            inService = attributes.get("name","")
            self.updateServiceEntry("service", inService)
            inProduct = attributes.get("product","")
            self.updateServiceEntry("product", inProduct)
            inExtraInfo = attributes.get("extrainfo","")
            self.updateServiceEntry("extrainfo", inExtraInfo)
            inConf = attributes.get("conf", "")
            self.updateServiceEntry("confidence", inConf)
        elif name == "state":
            inState = attributes.get("state", "")
            self.updateServiceEntry("state", inState)
        elif name == "osmatch":
            inOS = attributes.get("name","")
            self.updateHostEntry("osName", inOS)
        elif name == "script":
            inScriptID = attributes.get("id", "")
            inScriptOutput = attributes.get("output", "")
            self.insertNewScriptResult(inScriptID, inScriptOutput)
        elif name == "hostscript":
            self.inPortId = "Host Script"
            self.inProtocol = "Host Script"
        
    # adds a new host to the database
    def insertNewHost(self, ipAddr): 
        t = (ipAddr, ipAddr)
        hostSqlInsert = "INSERT INTO hosts (ipaddress) SELECT ? WHERE NOT EXISTS (SELECT 1 FROM hosts WHERE ipaddress = ?);"
        try:
            self.cursor.execute(hostSqlInsert,t)
            self.conn.commit()
        except sqlite3.OperationalError as strerror:
            logging.error ("SQLITE error - insertNewHost: " + str(strerror))  
            #sys.exit(2)

    #updates a host with relevant data to that host
    def updateHostEntry(self, attributeName, attributeValue):
        # update hosts set macaddress = '12:34:56:B1:D0:60' where (macaddress = '' or macaddress IS NULL) and (ipaddress = '192.168.43.1')
        t = (attributeValue, self.inIPaddr)
        hostSqlUpdate = "update hosts set "
        hostSqlUpdate += attributeName
        hostSqlUpdate += " = ? where ("
        hostSqlUpdate += attributeName 
        hostSqlUpdate += " = '' or "
        hostSqlUpdate += attributeName
        hostSqlUpdate += " IS NULL) and ("
        hostSqlUpdate += "ipaddress = ?)"
        
        try:
            self.cursor.execute(hostSqlUpdate, t)
            self.conn.commit()
        except sqlite3.OperationalError as strerror:
            logging.error ("SQLITE error - updateHostEntry: " + str(strerror))  
    
    # inserts a new service entry
    def insertNewService(self, ipAddr, proto, portid):
#        print (ipAddr + "  " + proto + "  " + portid + "  ") 
        t = (ipAddr, proto, portid, ipAddr, proto, portid)
        serviceSqlInsert = "INSERT INTO services (ipaddress, protocol, port) SELECT ?,?,? WHERE NOT EXISTS (SELECT 1 FROM services WHERE ipaddress = ? AND protocol = ? AND port = ?)"
        try:
            self.cursor.execute(serviceSqlInsert,t)
            self.conn.commit()
        except sqlite3.OperationalError as strerror:
            logging.error ("SQLITE error - insertNewService: " + str(strerror))  
        except sqlite3.IntegrityError as strerror:
            logging.error ("SQLITE error - insertNewService: " + str(strerror))  
        
    # update a particular service with relevant data 
    def updateServiceEntry(self, attributeName, attributeValue):
        t= (attributeValue,self.inIPaddr,self.inProtocol, self.inPortId)
        serviceSqlUpdate = "update services set " 
        serviceSqlUpdate += attributeName 
        serviceSqlUpdate += " = ? where ipaddress = ? and protocol = ? and port = ?"
        
        try:
            self.cursor.execute(serviceSqlUpdate,t)
            self.conn.commit()
        except sqlite3.OperationalError as strerror:
            logging.error ("SQLITE error - updateServiceEntry: " + str(strerror))     
        
    # inserts a new service entry
    def insertNewScriptResult(self, nmapScriptId, nmapScriptOutput):
        t = (self.inIPaddr,self.inProtocol, self.inPortId, nmapScriptId,nmapScriptOutput,self.inIPaddr,self.inProtocol,self.inPortId, nmapScriptId)
        scriptResultSqlInsert = "INSERT INTO scriptResults (ipaddress,protocol,port,nmapscriptid,scriptoutput) SELECT ?,?,?,?,? WHERE NOT EXISTS (SELECT 1 FROM scriptResults WHERE ipaddress = ? AND protocol = ? AND port = ? and nmapscriptid = ?)"
        try:
            self.cursor.execute(scriptResultSqlInsert,t)
            self.conn.commit()    
        except sqlite3.OperationalError as strerror:
            logging.error ("SQLITE error - insertNewScriptResult: " + str(strerror))                
            
    def printTable(self, tableName):
        self.cursor.execute('select * from ' + tableName)
        for row in self.cursor:
            logging.info (row)

